Add deterministic redistribute sharding for KafkaIO read. by tomstepp · Pull Request #36112 · apache/beam

tomstepp · 2025-09-10T18:42:04Z

Add deterministic redistribute sharding for KafkaIO read.

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

gemini-code-assist

Summary of Changes

Hello @tomstepp, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances the data redistribution mechanism within Apache Beam by introducing a deterministic sharding option. This feature is specifically integrated into the KafkaIO connector to provide predictable data distribution, which is crucial for scenarios requiring consistent processing, such as offset deduplication.

Highlights

Introduction of Deterministic Sharding: A new "deterministicSharding" option has been added to the RedistributeArbitrarily transform, allowing elements to be sharded based on their hash code rather than a random assignment.
KafkaIO Integration: The KafkaIO.Read transform now leverages this new deterministic sharding capability when offset deduplication is enabled, ensuring consistent distribution of Kafka records.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

tomstepp · 2025-09-10T18:45:49Z

R: @scwhittle

github-actions · 2025-09-10T18:46:55Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

tomstepp · 2025-09-12T16:14:50Z

Looks like the failures are unrelated and flaky tests (JmsIO and BeamFnLoggingClient)

tomstepp · 2025-09-23T01:04:38Z

I don't understand why PreCommits are failing now since they pass locally and the lines mentioned in errors don't line up with the current KafkaIOTest file.

liferoad · 2025-09-23T01:16:55Z

https://develocity.apache.org/s/kpt4uubaoc7b4

Compilation failed; see the compiler output below.
/runner/_work/beam/beam/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java:827: error: no suitable method found for mkKafkaReadTransform(int,int,ValueAsTimestampFn,boolean,boolean,<null>,<null>,<null>)
        mkKafkaReadTransform(
        ^
    method KafkaIOTest.mkKafkaReadTransform(int,@org.checkerframework.checker.nullness.qual.Nullable SerializableFunction<KV<Integer,Long>,Instant>) is not applicable
      (actual and formal argument lists differ in length)
    method KafkaIOTest.mkKafkaReadTransform(int,@org.checkerframework.checker.nullness.qual.Nullable Integer,@org.chec

tomstepp · 2025-09-23T01:33:20Z

https://develocity.apache.org/s/kpt4uubaoc7b4

Compilation failed; see the compiler output below.
/runner/_work/beam/beam/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java:827: error: no suitable method found for mkKafkaReadTransform(int,int,ValueAsTimestampFn,boolean,boolean,<null>,<null>,<null>)
        mkKafkaReadTransform(
        ^
    method KafkaIOTest.mkKafkaReadTransform(int,@org.checkerframework.checker.nullness.qual.Nullable SerializableFunction<KV<Integer,Long>,Instant>) is not applicable
      (actual and formal argument lists differ in length)
    method KafkaIOTest.mkKafkaReadTransform(int,@org.checkerframework.checker.nullness.qual.Nullable Integer,@org.chec

However I don't see mkKafkaReadTransform at KafkaIOTest.java:827 (which is a blank line currently)

liferoad · 2025-09-23T02:17:57Z

try to change null /*topics*/, to null, /*topics*/

tomstepp · 2025-09-23T04:28:29Z

try to change null /*topics*/, to null, /*topics*/

Fixed this but the error message still exists and doesn't really make sense to me, not sure what I'm missing.

…key.

liferoad · 2025-09-23T17:00:26Z

./gradlew :sdks:java:io:kafka:compileTestJava: this runs fine with your PR. Not sure what's going on.

tomstepp · 2025-09-23T17:17:52Z

./gradlew :sdks:java:io:kafka:compileTestJava: this runs fine with your PR. Not sure what's going on.

@liferoad, I realized that the line mismatches were probably due to pulling master into the branch for testing. After pulling master, resolving conflicts, and pushing, it looks like this is resolved now. Thanks for your help!

The latest errors I think are flakes or unrelated to this change.

codecov · 2025-09-24T17:58:47Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 56.79%. Comparing base (0384a59) to head (06469c5).
⚠️ Report is 25 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #36112      +/-   ##
============================================
+ Coverage     54.92%   56.79%   +1.86%     
- Complexity     1617     3386    +1769     
============================================
  Files          1057     1220     +163     
  Lines        164391   185653   +21262     
  Branches       1165     3523    +2358     
============================================
+ Hits          90289   105435   +15146     
- Misses        71953    76884    +4931     
- Partials       2149     3334    +1185

Flag	Coverage Δ
java	`70.39% <ø> (+2.03%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

scwhittle

Just some minor comments. I'm not super familiar with the translation stuff, it might be a good idea to have someone else review for that.

…ify sharding Fn display name.

tomstepp · 2025-09-25T15:08:39Z

R: @chamikaramj

tomstepp · 2025-09-25T15:09:47Z

Just some minor comments. I'm not super familiar with the translation stuff, it might be a good idea to have someone else review for that.

Thanks! Added Cham for translation.

tomstepp · 2025-09-29T17:01:29Z

@chamikaramj could you please review?

tomstepp · 2025-09-30T17:37:00Z

I think the non-translation LG to Sam and translation LG to Cham, so should be ready to merge by either.

chamikaramj · 2025-09-30T18:53:03Z

Thanks. Re-ran the failing tests.

github-actions Bot added java io kafka labels Sep 10, 2025

gemini-code-assist Bot reviewed Sep 10, 2025

View reviewed changes

tomstepp commented Sep 10, 2025

View reviewed changes

Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java Outdated

scwhittle requested changes Sep 11, 2025

View reviewed changes

Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java

Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java Outdated

Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java Outdated

tomstepp requested a review from scwhittle September 11, 2025 16:22

scwhittle requested changes Sep 16, 2025

View reviewed changes

Comment thread sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java

Comment thread sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/Redistribute.java Outdated

scwhittle requested changes Sep 17, 2025

View reviewed changes

stankiewicz requested changes Sep 17, 2025

View reviewed changes

Comment thread sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java Outdated

tomstepp requested review from scwhittle and stankiewicz September 19, 2025 19:00

stankiewicz reviewed Sep 22, 2025

View reviewed changes

Comment thread sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java Outdated

stankiewicz reviewed Sep 22, 2025

View reviewed changes

Comment thread sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java

scwhittle requested changes Sep 22, 2025

View reviewed changes

Comment thread sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java

Comment thread sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaReadRedistribute.java Outdated

tomstepp requested review from scwhittle and stankiewicz September 22, 2025 22:42

tomstepp added 5 commits September 23, 2025 08:25

Add deterministic redistribute sharding for KafkaIO read.

403e2b4

Address PR feedback.

a84ab49

Provide more detailed transform name for the redistribute.

0627d94

Address spotless precommit findings.

bc96165

Address spotless precommit findings.

65621c8

tomstepp added 11 commits September 23, 2025 08:25

Keep redistribute transform name the same.

a611d9b

Add deterministic sharding unit test.

1ba7206

Refactor to specific deterministic Kafka redistribute method.

b4abac1

Add redistribute by key variant.

b36e2ec

Add test of sharding fns.

4496030

Add bucketing to redistributeByKey and add option to redistribute by …

99d8b78

…key.

Actually enable withRedistributeByRecordKey in KafkaIOTest.

89a1244

Add byRecordKey property to Kafka read compatibility.

760f297

Fix comma formatting for mkKafkaReadTransform.

638feab

Fix cases where reader was not overwritten when building Kafka reader.

1447890

Rebase and revert method rename for debugging.

6489053

tomstepp force-pushed the kafka-key branch from da5ed26 to 6489053 Compare September 23, 2025 15:37

Address spotless finding for makeKafkaRecord.

cf907b7

Add tests for deterministic sharding.

06469c5

scwhittle requested changes Sep 25, 2025

View reviewed changes

numBuckets as UnsignedInteger to reduce conversion overhead, and clar…

ebef848

…ify sharding Fn display name.

chamikaramj reviewed Sep 30, 2025

View reviewed changes

Comment thread sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java

tomstepp requested a review from scwhittle September 30, 2025 17:09

stankiewicz approved these changes Sep 30, 2025

View reviewed changes

chamikaramj merged commit 4d87e04 into apache:master Sep 30, 2025
23 of 28 checks passed

Uh oh!

Conversation

tomstepp commented Sep 10, 2025

GitHub Actions Tests Status (on master branch)

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

Uh oh!

tomstepp commented Sep 10, 2025

Uh oh!

github-actions Bot commented Sep 10, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tomstepp commented Sep 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tomstepp commented Sep 23, 2025

Uh oh!

liferoad commented Sep 23, 2025

Uh oh!

tomstepp commented Sep 23, 2025

Uh oh!

liferoad commented Sep 23, 2025

Uh oh!

tomstepp commented Sep 23, 2025

Uh oh!

liferoad commented Sep 23, 2025

Uh oh!

tomstepp commented Sep 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov Bot commented Sep 24, 2025

Codecov Report

Uh oh!

scwhittle left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tomstepp commented Sep 25, 2025

Uh oh!

tomstepp commented Sep 25, 2025

Uh oh!

tomstepp commented Sep 29, 2025

Uh oh!

Uh oh!

tomstepp commented Sep 30, 2025

Uh oh!

chamikaramj commented Sep 30, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

tomstepp commented Sep 12, 2025 •

edited

Loading

tomstepp commented Sep 23, 2025 •

edited

Loading